#ifndef SILL_LEARNING_DMOVE_MOVE_LEAF_HPP
#define SILL_LEARNING_DMOVE_MOVE_LEAF_HPP

#include <set>
#include <sill/iterator/subset_iterator.hpp>
#include <sill/learning/structure_old/decomposable_move.hpp>
#include <sill/base/stl_util.hpp>
#include <sill/macros_def.hpp>

namespace sill {

  /**
   * Class for representing a possible move in structure search for
   * decomposable models:
   *  Move all variables which only appear in some leaf to a newly created leaf.
   *  The old leaf is removed. The new leaf must share at least 1 variable with
   *  whatever node it is attached to.
   *
   * For maximal and non-maximal JTs.
   *
   * @param F         type of factor used in the decomposable model
   * @param Inserter  Inserter.push(m, score) must insert move m with
   *                  score/estimate score into whatever container collects
   *                  new moves which are generated by this class.
   *
   * \author Joseph Bradley
   * \ingroup learning_structure
   */
  template <typename F, typename Inserter>
  class dmove_move_leaf
    : public decomposable_move<F,Inserter> {

  public:
    //! The base (move) type
    typedef decomposable_move<F,Inserter> base;

    //! The type of factor used in the decomposable model
    typedef F factor_type;

    //! The type of variable associated with a factor
    typedef typename F::variable_type variable_type;

    //! The domain type of the factor
    typedef typename base::domain_type domain_type;

    //! The generic representation for clique changes
    typedef typename base::clique_change clique_change;

    //! The type of edge associated with the model
    typedef typename base::edge edge;

    //! The type of vertex associated with the model
    typedef typename base::vertex vertex;

    //////////////////////// PRIVATE DATA AND METHODS ///////////////////////

  private:
    //! Leaf to move
    vertex move_v;
    //! Vertex to attach move_v to
    vertex attach_v;
    //! Variables in attach_v to add to the new leaf formed by move_v
    domain_type new_vars;

    ////////////////////////// PUBLIC METHODS //////////////////////////////
  public:
    //! Creates a null move (which can be used to generate new moves).
    dmove_move_leaf() : move_v(decomposable<F>::null_vertex()),
                        attach_v(decomposable<F>::null_vertex()) { }

    //! Creates a move which moves leaf move_v to a new leaf attached to
    //! vertex attach_v.
    dmove_move_leaf(vertex move_v, vertex attach_v, const domain_type& new_vars)
      : move_v(move_v), attach_v(attach_v), new_vars(new_vars) { }

    //! Given a model, score, and statistics class, generate all possible moves
    //! and insert pointers to them (and their scores) into the provided
    //! Inserter.
    void
    generate_all_moves(const learnt_decomposable<F>& model, double cur_score,
                       const decomposable_score<F>& score, dataset_statistics<>& stats,
                       Inserter& inserter, bool use_estimates = false) const {
      // For each (leaf node, other node) pair,
      foreach(vertex u, model.vertices()) {
        if (model.out_degree(u) > 1)
          continue;
        vertex neighbor(decomposable<F>::null_vertex());
        foreach(vertex n, model.neighbors(u)) {
          neighbor = n;
          break;
        }
        domain_type move_vars(set_difference(model.clique(u), model.clique(neighbor)));
        size_t max_size_new_vars(model.max_clique_size() - move_vars.size());
        foreach(vertex v, model.vertices()) {
          if (u == v || v == neighbor)
            continue;
          // For each possible set of new_vars,
          subset_iterator<domain_type> new_vars_it, new_vars_it_end;
          if (model.maximal_JTs())
            new_vars_it = subset_iterator<domain_type>(model.clique(v),
                                                          max_size_new_vars);
          else
            new_vars_it = subset_iterator<domain_type>(model.clique(v),
                                                          1,max_size_new_vars);
          while(new_vars_it != new_vars_it_end) {
            // Generate the corresponding new move.
            dmove_move_leaf<F,Inserter>* move_ptr =
              new dmove_move_leaf<F,Inserter>(u, v, *new_vars_it);
            double change;
            if (use_estimates)
              change = score.estimate_change(model, cur_score,
                                             *move_ptr, stats).second;
            else
              change = score.compute_change(model, cur_score,
                                            *move_ptr, stats).second;
            inserter.push(move_ptr, change);
            ++new_vars_it;
          }
        }
      }
    }

    //! Given a model, score, another decomposable_move (which has just been
    //! committed), and statistics class,
    //! generate all possible new moves of this type and insert pointers to
    //! them (and their scores) into the provided Inserter.
    void
    generate_new_moves(const learnt_decomposable<F>& model, double cur_score,
                       const decomposable_score<F>& score,
                       const std::vector<clique_change>& clique_changes,
                       dataset_statistics<>& stats, Inserter& inserter,
                       bool use_estimates = false) const {
      // Pairs (changed leaf nodes, other nodes)
      std::set<vertex> changed_leaf_nodes;
      foreach(const clique_change& change, clique_changes) {
        if (model.in_degree(change.v) > 1)
          continue;
        changed_leaf_nodes.insert(change.v);
        vertex neighbor(decomposable<F>::null_vertex());
        foreach(vertex n, model.neighbors(change.v)) {
          neighbor = n;
          break;
        }
        size_t max_size_new_vars
          (model.max_clique_size()
           - set_difference(model.clique(change.v), model.clique(neighbor)).size());
        foreach(vertex u, model.vertices()) {
          if (change.v == u || u == neighbor)
            continue;
          // For each new possible set new_vars
          subset_iterator<domain_type> new_vars_it, new_vars_it_end;
          if (model.maximal_JTs())
            new_vars_it = subset_iterator<domain_type>(model.clique(u),
                                                          max_size_new_vars);
          else
            new_vars_it = subset_iterator<domain_type>(model.clique(u),
                                                          1, max_size_new_vars);
          while(new_vars_it != new_vars_it_end) {
            // Generate the corresponding new move.
            dmove_move_leaf<F,Inserter>* move_ptr =
              new dmove_move_leaf<F,Inserter>(change.v, u, *new_vars_it);
            double change;
            if (use_estimates)
              change = score.estimate_change(model, cur_score,
                                             *move_ptr, stats).second;
            else
              change = score.compute_change(model, cur_score,
                                            *move_ptr, stats).second;
            inserter.push(move_ptr, change);
            ++new_vars_it;
          }
        }
      }
      // Pairs (unchanged leaf nodes, changed leaf nodes)
      foreach(vertex u, model.vertices()) {
        if (model.in_degree(u) > 1 || changed_leaf_nodes.count(u))
          continue;
        vertex neighbor(decomposable<F>::null_vertex());
        foreach(vertex n, model.neighbors(u)) {
          neighbor = n;
          break;
        }
        size_t max_size_new_vars(model.max_clique_size());
        if (neighbor != decomposable<F>::null_vertex())
          max_size_new_vars -=
            set_difference(model.clique(u), model.clique(neighbor)).size();
        foreach(const clique_change& change, clique_changes) {
          if (model.in_degree(change.v) > 1)
            continue;
          // For each new possible set new_vars
          subset_iterator<domain_type> new_vars_it, new_vars_it_end;
          if (model.maximal_JTs())
            new_vars_it =
              subset_iterator<domain_type>(model.clique(change.v),
                                              max_size_new_vars);
          else
            new_vars_it =
              subset_iterator<domain_type>(model.clique(change.v),
                                              1, max_size_new_vars);
          while(new_vars_it != new_vars_it_end) {
            ++new_vars_it;
            // Generate the corresponding new move.
            dmove_move_leaf<F,Inserter>* move_ptr =
              new dmove_move_leaf<F,Inserter>(u, change.v, *new_vars_it);
            double change;
            if (use_estimates)
              change = score.estimate_change(model, cur_score,
                                             *move_ptr, stats).second;
            else
              change = score.compute_change(model, cur_score,
                                            *move_ptr, stats).second;
            inserter.push(move_ptr, change);
          }
        }
      }
    }

    //! Call a score functor on each clique/separator inserted/deleted
    //! by this move (where the move is represented by clique/separator
    //! insertions/deletions).
    //! @return false if move is invalid, else true
    bool map_score_functor(decomposable_score_functor<F>& func,
                           const learnt_decomposable<F>& model,
                           dataset_statistics<>& stats) const {
      // Check the validity
      if (!(model.contains(move_v)) || !(model.contains(attach_v)) ||
          model.out_degree(move_v) > 1 ||
          !(includes(model.clique(attach_v), new_vars)))
        return false;
      domain_type new_clique(model.clique(move_v));
      vertex neighbor(decomposable<F>::null_vertex());
      foreach(vertex u, model.neighbors(move_v)) {
        neighbor = u;
        new_clique = set_difference(new_clique, model.clique(u));
        break;
      }
      if (new_clique.size() == 0)
        return false;
      new_clique = set_union(new_clique, new_vars);
      if (model.maximal_JTs()) {
        if (new_clique.size() != model.max_clique_size())
          return false;
      } else {
        if (new_clique.size() > model.max_clique_size())
          return false;
      }
      // Call functor on changes
      func.deleted_clique(model.marginal(move_v));
      const F& new_factor = stats.marginal(new_clique);
      func.inserted_clique(new_factor);
      if (neighbor != decomposable<F>::null_vertex())
        func.deleted_separator(model.marginal(model.get_edge(move_v,
                                                             neighbor)));
      func.inserted_separator(new_factor.marginal(new_vars));
      return true;
    }

    //! Given a model, check to see if the move is still valid.
    bool valid(const learnt_decomposable<F>& model) const {
      if (!(model.contains(move_v)) || !(model.contains(attach_v)) ||
          model.out_degree(move_v) > 1 ||
          !(includes(model.clique(attach_v), new_vars)))
        return false;
      domain_type new_clique(model.clique(move_v));
      foreach(vertex u, model.neighbors(move_v)) {
        new_clique = set_difference(new_clique, model.clique(u));
        break;
      }
      if (new_clique.size() == 0)
        return false;
      new_clique = set_union(new_clique, new_vars);
      if (model.maximal_JTs()) {
        if (new_clique.size() != model.max_clique_size())
          return false;
      } else {
        if (new_clique.size() > model.max_clique_size())
          return false;
      }
      return true;
    }

    //! Given a mutable model, commit the move.
    //! Note: This does NOT check if the move is valid!
    //! Also, this does not calibrate or renormalize the model.
    std::vector<clique_change>
    commit(learnt_decomposable<F>& model, dataset_statistics<>& stats) const {
      std::vector<clique_change> changes;
      domain_type new_clique(model.clique(move_v));
      vertex neighbor(decomposable<F>::null_vertex());
      domain_type deleted_vars;
      domain_type added_vars(new_vars);
      foreach(vertex u, model.neighbors(move_v)) {
        neighbor = u;
        deleted_vars = set_intersect(new_clique, model.clique(u));
        new_clique = set_difference(new_clique, deleted_vars);
        break;
      }
      new_clique = set_union(new_clique, new_vars);
      if (neighbor != decomposable<F>::null_vertex())
        model.remove_edge(move_v, neighbor);
      model.set_clique(move_v, new_clique, stats.marginal(new_clique));
      model.add_edge(move_v, attach_v);
      changes.push_back(clique_change(move_v, deleted_vars, added_vars));
      return changes;
    }

  }; // class dmove_move_leaf

} // namespace sill

#include <sill/macros_undef.hpp>

#endif // #ifndef SILL_LEARNING_DMOVE_MOVE_LEAF_HPP
